import json
from operator import itemgetter
from typing import Any, Dict, List, Optional, TypedDict

from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool, render_text_description
from langchain_core.runnables import RunnableConfig
from langchain_core.utils.function_calling import convert_to_openai_tool

from ChatGLM_new import tongyi_llm,zhipu_llm,xinghuo_llm,local_llm,sspy_llm,ollama_llm,tongyi_llm_openai,kimi_llm,kimi_llm_openai,one_api_llm,xinghuo_llm_openai
from langchain_core.pydantic_v1 import BaseModel, Field

class RoadModel(BaseModel):
    """Information about a road."""
    roadName: Optional[str] = Field(default=None, description="道路名称")
    roadSide: Optional[str] = Field(default=None, description="道路路侧")
    dimmingValue: Optional[int] = Field(default=None, description="调光值,开灯时默认为100,关灯时默认为0", ge=0, le=100)

class LamppostModel(BaseModel):
    """Information about pole."""
    lamppostCode: Optional[str] = Field(default=None, description="灯杆编号")
    dimmingValue: Optional[int] = Field(default=None, description="调光值,开灯时默认为100,关灯时默认为0", ge=0, le=100)

class PlanModel(BaseModel):
    """方案的相关信息."""
    name: Optional[str] = Field(default=None, description="方案名称，必须由用户指定方案名称")


@tool
def plan_promote(planModelList: List[PlanModel]):
    """ 按方案进行开关灯，方案的名称必须由用户指定，用户未指定时，请询问用户方案的名称信息。参数格式为json:{"planModelList": [{"name": "..."}]}"""
    init_prompt = """
    缺少方案名称，请询问用户更多信息，引导用户输入需要下发的方案名称
    """
    try:
        for planModel in planModelList:
            if planModel.name is None or planModel.name == "" :
                return init_prompt
            return "方案名称不存在"
    except:
        return init_prompt
    return init_prompt

@tool
def road_promote(roadModelList: List[RoadModel]):
    """ 按道路名称进行开关灯。参数格式为json:{"roadModelList": [{"roadName": "舜华路", "roadSide": "", "dimmingValue": 100},{"roadName": "和平路", "roadSide": "", "dimmingValue": 0}]}"""
    init_prompt = """
    缺少道路名称或者调光值，请询问用户更多信息，引导用户输入道路名称和调光值
    """
    try:
        for roadModel in roadModelList:
            if roadModel.roadName is None or roadModel.roadName == "" or roadModel.dimmingValue is None or roadModel.dimmingValue == "":
                return init_prompt
            return "操作成功"
    except:
        return init_prompt
    return init_prompt

@tool
def pole_promote(lamppostModelList: List[LamppostModel]):
    """按灯杆编号进行开关灯。参数格式为json:{"lamppostModelList": [{"lamppostCode": "201", "dimmingValue": 100},{"lamppostCode": "202407051130", "dimmingValue": 100}]}"""
    init_prompt = """
        缺少灯杆编号或者调光值，请询问用户更多信息，引导用户输入灯杆编号和调光值
        """
    try:
        for lamppostModel in lamppostModelList:
            if lamppostModel.lamppostCode is None or lamppostModel.dimmingValue is None:
                return init_prompt
            return "操作失败"
    except:
        return init_prompt

@tool
def ask_for_more_info(question):
    """询问用户更多信息,对用户进行提问。参数question是字符串类型,格式：{"question": ".....？"} """
    #print(f"问题：{question}")
    return input("请输入更多信息：")

tools = [road_promote, pole_promote, plan_promote, ask_for_more_info]
rendered_tools = render_text_description(tools)
openai_tools = [convert_to_openai_tool(tool) for tool in tools]
#print(openai_tools)



# system_prompt = f"""您是一名助理，可以使用以下工具集。 以下是每个工具的名称和说明:
#
# {rendered_tools}
#
# 根据用户输入，返回要使用的工具的名称和输入。 将您的响应作为带有'name'和'arguments'键的 JSON blob 返回，“arguments”键对应的值应该是所选函数的输入参数的字典，字典里不要有任何说明,此JSON blob必须是如下格式：```json
# ...
# ```"""
# prompt = ChatPromptTemplate.from_messages(
#     [("system", system_prompt), ("user", "{input}")]
# )
#
# chain = prompt | tongyi_llm | JsonOutputParser()
# print(chain.invoke({"input": "我要开灯"}))


tool_names = [i.name for i in tools]
parser = JsonOutputParser()
def tool_chain(model_output):
    tool_map = {tool.name: tool for tool in tools}
    chosen_tool = tool_map[model_output["name"]]
    return itemgetter("arguments") | chosen_tool

def get_args(res):
    args=res.split("Action:")[1].split("Action Input:")
    try:
        arguments = parser.parse(args[1].strip())
    except:
        arguments = {}
    return {"name": args[0].strip(), "arguments": arguments}

def prompt_react(input):
    PREFIX = f"""你是一个开关灯助手，只能回答开关灯相关的问题，如果用户提出的问题与开关灯无关，需要引导用户提问开关灯相关的问题。
    不要进行任何猜测和假设,如果需要用户回答问题,请将问题引导到ask_for_more_info
    不要进行任何猜测和假设,如果需要用户回答问题,请将问题引导到ask_for_more_info
    不要进行任何猜测和假设,如果需要用户回答问题,请将问题引导到ask_for_more_info
    
    不要有任何说明、解释、注释
    不要有任何说明、解释、注释
    不要有任何说明、解释、注释
    
    有以下工具可以调用：
        {openai_tools}
    请尽可能从用户输入的问题中找到合适的工具
    
    "Final Answer的作用只输出最终的答案,如果需要用户回答问题,请将问题引导到ask_for_more_info"
    "Final Answer的作用只输出最终的答案,如果需要用户回答问题,请将问题引导到ask_for_more_info"
    """
    prompt = f"""
        Use the following format:
        Question: 用户的信息
        Thought: 你应该经常考虑该怎么做
        Action: 按顺序采取的行动应为以下操作之一：{tool_names}
        Action Input: 操作的输入，不要有任何解释和说明，直接输入，不能包含注释，不能包含注释，不能包含注释，不能包含注释，不能包含注释，不能有"//"
        Observation: 用户最新的输入信息
        ... (this Thought/Action/Action Input/Observation can repeat N times)
        Thought: 我现在知道了最终的答案
        Final Answer: 问题的最终答案。需要用户输入信息时，不要输出这个最终答案
        
        Begin!
        
        Question: {input}
        """
    return [{"role": "system", "content": PREFIX},{"role": "user", "content": prompt}]

#chain = prompt_react | tongyi_llm
query = input("请输入问题：")
p_r = prompt_react(query)
while True:
    print(p_r[1]["content"])
    res = xinghuo_llm_openai.invoke(p_r,stop=["\nObservation:"])
    res_content = res.content
    print("\033[1;32;40m" + res_content + "\033[0m")
    if res_content.find("Action Input:") != -1:
        args = get_args(res_content)
        t_run = tool_chain(args)
        if args["name"] == "ask_for_more_info":
            print("\033[1;31;40m" + args["arguments"]["question"] + "\033[0m")
            query = input("请输入问题：")
        else:
            query = t_run.invoke(args)
        p_r[1]["content"] = p_r[1]["content"] + res_content + "\nObservation: " + str(query) + "\n"
    elif res_content.find("Final Answer:") != -1:
        print("\033[1;31;40m" + res_content.split("Final Answer:")[1] + "\033[0m")
        query = input("请输入问题：")
        p_r = prompt_react(query)
    else:
        p_r[1]["content"] = p_r[1]["content"] + res_content
        break